package ai.mspbots.poc.help;

import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.embeddings.wordvectors.WordVectors;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class SearchHelper {

    private final RestHighLevelClient client;

    public SearchHelper(RestHighLevelClient client) {
        this.client = client;
    }

    private static final String indexName = "sys_report_vector";

    public SearchResponse searchByText(String vectorField, String text) {
        return searchByVector(indexName, vectorField, text);
    }

    public SearchResponse searchByVector(String index, String vectorField, String text) {
        // 创建搜索请求
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 准备脚本参数
        Map<String, Object> params = new HashMap<>();
        params.put("query_vector", VectorHelper.vectorizedText(text));

        // 使用 script_score 查询
        Script script = new Script(
                ScriptType.INLINE,
                "painless",
                "cosineSimilarity(params.query_vector, doc['" + vectorField + "']) + 1.0",
                params
        );

        sourceBuilder.query(QueryBuilders.scriptScoreQuery(QueryBuilders.matchAllQuery(), script));

        sourceBuilder.size(10); // 设置返回的结果数量
        searchRequest.source(sourceBuilder);

        // 执行搜索
        SearchResponse searchResponse = null;
        try {
            searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        // System.out.println("搜索结果: " + searchResponse.toString());
        return searchResponse;
    }

    public void searchByVector(String index, String vectorField, double[] queryVector) throws IOException {
        // 创建搜索请求
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 准备脚本参数
        Map<String, Object> params = new HashMap<>();
        params.put("query_vector", queryVector);

        // 使用 script_score 查询
        Script script = new Script(
                ScriptType.INLINE,
                "painless",
                "cosineSimilarity(params.query_vector, doc['" + vectorField + "']) + 1.0",
                params
        );

        sourceBuilder.query(QueryBuilders.scriptScoreQuery(QueryBuilders.matchAllQuery(), script));

        sourceBuilder.size(10); // 设置返回的结果数量
        searchRequest.source(sourceBuilder);

        // 执行搜索
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        // System.out.println("搜索结果: " + searchResponse.toString());
    }


}
